from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
from datasets import Dataset, Audio
from typing import Union, Iterable
import torch
import logging

logger = logging.getLogger(__name__)

class ImprovedWhisperASR:
    def __init__(self,
                 device='cuda:9',
                 max_new_tokens=1024,
                 language='chinese',
                 task='transcribe',
                 chunk_length_s=30,
                 batch_size=16,
                 word_level_timestamps=True):
        """
        改进的Whisper ASR实现，基于参考实现的最佳实践

        Args:
            device: 指定GPU设备
            max_new_tokens: 最大生成token数
            language: 转录语言
            task: 任务类型
            chunk_length_s: 音频分块长度（秒）
            batch_size: 批处理大小
            word_level_timestamps: 是否启用词级时间戳
        """
        self.device = device if torch.cuda.is_available() else 'cpu'
        logger.info(f"🎤 Whisper ASR 初始化，使用设备: {self.device}")

        # 配置数据类型
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

        # 模型ID
        model_id = 'openai/whisper-large-v3'
        logger.info(f"📥 正在加载模型: {model_id}")

        # 加载模型和处理器
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
            use_safetensors=True
        )
        model.to(self.device)

        processor = AutoProcessor.from_pretrained(model_id)

        # 创建pipeline
        self.pipe = pipeline(
            'automatic-speech-recognition',
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            chunk_length_s=chunk_length_s,
            max_new_tokens=max_new_tokens,
            batch_size=batch_size,
            return_timestamps=True if not word_level_timestamps else 'word',
            torch_dtype=torch_dtype,
            device=self.device
        )

        self._kwargs = {'language': language, 'task': task}
        logger.info(f"✅ Whisper ASR 初始化完成，词级时间戳: {'启用' if word_level_timestamps else '禁用'}")

    def transcribe(self, audio_file: Union[str, Iterable]) -> dict:
        """
        转录音频文件

        Args:
            audio_file: 音频文件路径或文件列表

        Returns:
            包含转录结果和时间戳的字典
        """
        logger.info(f"🎵 开始转录音频: {audio_file}")

        # 创建数据集
        ds = Dataset.from_dict({'audio': [audio_file] if isinstance(audio_file, str) else audio_file})\
            .cast_column('audio', Audio(sampling_rate=16000))

        # 处理音频
        result = self._handle(ds)

        if isinstance(audio_file, str):
            return result[0]
        else:
            return result

    def _handle(self, ds: Dataset) -> list:
        """处理数据集"""
        results = []

        for data in ds:
            sample = data['audio']
            logger.info(f"🔊 处理音频片段，时长: {sample['array'].size / sample['sampling_rate']:.2f}秒")

            # 进行转录
            prediction = self.pipe(sample.copy(), generate_kwargs=self._kwargs)

            # 解析结果
            if 'chunks' in prediction:
                # 处理带时间戳的结果
                chunks = prediction['chunks']
                processed_chunks = []

                for chunk in chunks:
                    if isinstance(chunk, dict) and 'text' in chunk:
                        # 安全处理时间戳
                        timestamp = chunk.get('timestamp', [0, 0])
                        if isinstance(timestamp, (list, tuple)) and len(timestamp) >= 2:
                            start_time = float(timestamp[0])
                            end_time = float(timestamp[1])
                        else:
                            start_time = 0.0
                            end_time = 0.0

                        processed_chunk = {
                            'text': chunk['text'].strip(),
                            'start': start_time,
                            'end': end_time
                        }

                        # 如果有词级时间戳，添加到结果中
                        if 'words' in chunk:
                            processed_chunk['words'] = chunk['words']

                        if processed_chunk['text']:  # 只添加非空文本
                            processed_chunks.append(processed_chunk)

                results.append({
                    'text': prediction['text'].strip(),
                    'chunks': processed_chunks
                })
            else:
                # 处理简单文本结果
                results.append({
                    'text': prediction['text'].strip(),
                    'chunks': []
                })

        logger.info(f"✅ 转录完成，共处理 {len(results)} 个音频片段")
        return results

    def extract_word_timestamps(self, result: dict) -> list:
        """
        从转录结果中提取词级时间戳

        Args:
            result: transcribe方法返回的结果

        Returns:
            词级时间戳列表
        """
        word_timestamps = []

        if 'chunks' in result:
            for chunk in result['chunks']:
                if 'words' in chunk:
                    for word_info in chunk['words']:
                        if isinstance(word_info, dict):
                            word_timestamps.append({
                                'word': word_info.get('word', ''),
                                'start': word_info.get('start', 0),
                                'end': word_info.get('end', 0)
                            })
                else:
                    # 如果没有词级时间戳，将整个chunk作为一个词
                    if chunk['text']:
                        words = chunk['text'].split()
                        words_count = len(words)
                        if words_count > 0:
                            duration = chunk['end'] - chunk['start']
                            word_duration = duration / words_count

                            for i, word in enumerate(words):
                                start_time = chunk['start'] + i * word_duration
                                end_time = start_time + word_duration
                                word_timestamps.append({
                                    'word': word,
                                    'start': start_time,
                                    'end': end_time
                                })

        return word_timestamps